msg_tool\scripts\emote/
pimg.rs

1//! Emote Multiple Image File (.pimg)
2use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::try_option;
6use crate::types::*;
7use crate::utils::img::*;
8use crate::utils::psd::*;
9use crate::utils::struct_pack::*;
10use anyhow::Result;
11use emote_psb::PsbReader;
12use libtlg_rs::*;
13use std::collections::HashMap;
14use std::io::{Read, Seek};
15use std::path::Path;
16
17#[derive(Debug)]
18/// Emote PImg Script Builder
19pub struct PImgBuilder {}
20
21impl PImgBuilder {
22    /// Creates a new instance of `PImgBuilder`
23    pub const fn new() -> Self {
24        Self {}
25    }
26}
27
28impl ScriptBuilder for PImgBuilder {
29    fn default_encoding(&self) -> Encoding {
30        Encoding::Utf8
31    }
32
33    fn build_script(
34        &self,
35        buf: Vec<u8>,
36        filename: &str,
37        _encoding: Encoding,
38        _archive_encoding: Encoding,
39        config: &ExtraConfig,
40        _archive: Option<&Box<dyn Script>>,
41    ) -> Result<Box<dyn Script>> {
42        Ok(Box::new(PImg::new(MemReader::new(buf), filename, config)?))
43    }
44
45    fn build_script_from_file(
46        &self,
47        filename: &str,
48        _encoding: Encoding,
49        _archive_encoding: Encoding,
50        config: &ExtraConfig,
51        _archive: Option<&Box<dyn Script>>,
52    ) -> Result<Box<dyn Script>> {
53        if filename == "-" {
54            let data = crate::utils::files::read_file(filename)?;
55            Ok(Box::new(PImg::new(MemReader::new(data), filename, config)?))
56        } else {
57            let f = std::fs::File::open(filename)?;
58            let reader = std::io::BufReader::new(f);
59            Ok(Box::new(PImg::new(reader, filename, config)?))
60        }
61    }
62
63    fn build_script_from_reader(
64        &self,
65        reader: Box<dyn ReadSeek>,
66        filename: &str,
67        _encoding: Encoding,
68        _archive_encoding: Encoding,
69        config: &ExtraConfig,
70        _archive: Option<&Box<dyn Script>>,
71    ) -> Result<Box<dyn Script>> {
72        Ok(Box::new(PImg::new(reader, filename, config)?))
73    }
74
75    fn extensions(&self) -> &'static [&'static str] {
76        &["pimg"]
77    }
78
79    fn script_type(&self) -> &'static ScriptType {
80        &ScriptType::EmotePimg
81    }
82
83    fn is_this_format(&self, filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
84        if Path::new(filename)
85            .extension()
86            .map(|ext| ext.to_ascii_lowercase() == "pimg")
87            .unwrap_or(false)
88            && buf_len >= 4
89            && buf.starts_with(b"PSB\0")
90        {
91            return Some(255);
92        }
93        None
94    }
95
96    fn is_image(&self) -> bool {
97        true
98    }
99}
100
101struct PImgLayer<'a> {
102    data: &'a PsbValueFixed,
103    name: &'a str,
104    layer_id: i64,
105    /// seems is layer type in PSD files
106    layer_type: i64,
107    left: u32,
108    top: u32,
109    width: u32,
110    height: u32,
111    opacity: u8,
112    visible: bool,
113    type_: i64,
114    children: Vec<PImgLayer<'a>>,
115}
116
117impl std::fmt::Debug for PImgLayer<'_> {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        f.debug_struct("PImgLayer")
120            .field("layer_id", &self.layer_id)
121            .field("layer_type", &self.layer_type)
122            .field("name", &self.name)
123            .field("left", &self.left)
124            .field("top", &self.top)
125            .field("width", &self.width)
126            .field("height", &self.height)
127            .field("opacity", &self.opacity)
128            .field("visible", &self.visible)
129            .field("type", &self.type_)
130            .field("children", &self.children)
131            .finish()
132    }
133}
134
135impl<'a> PImgLayer<'a> {
136    pub fn new(data: &'a PsbValueFixed, layers: &'a PsbValueFixed) -> Result<Self> {
137        let layer_id = data["layer_id"]
138            .as_i64()
139            .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid layer_id"))?;
140        let layer_type = data["layer_type"]
141            .as_i64()
142            .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid layer_type"))?;
143        let name = data["name"]
144            .as_str()
145            .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid name"))?;
146        let left = data["left"].as_u32();
147        let top = data["top"].as_u32();
148        let width = data["width"].as_u32();
149        let height = data["height"].as_u32();
150        let (left, top, width, height) = if layer_type != 0 {
151            (
152                left.unwrap_or(0),
153                top.unwrap_or(0),
154                width.unwrap_or(0),
155                height.unwrap_or(0),
156            )
157        } else {
158            (
159                left.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid left"))?,
160                top.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid top"))?,
161                width.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid width"))?,
162                height.ok_or_else(|| anyhow::anyhow!("Layer does not have a valid height"))?,
163            )
164        };
165        let opacity = data["opacity"]
166            .as_u8()
167            .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid opacity"))?;
168        let visible = data["visible"].as_i64().unwrap_or(1) != 0;
169        let type_ = data["type"]
170            .as_i64()
171            .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid type"))?;
172        let mut children = Vec::new();
173        for layer in layers.members() {
174            if layer_type == 2 || layer_type == 1 {
175                if let Some(parent_id) = layer["group_layer_id"].as_i64() {
176                    if parent_id == layer_id {
177                        children.push(PImgLayer::new(layer, layers)?);
178                    }
179                }
180            } else if layer_type == 0 {
181                if let Some(base_id) = layer["diff_id"].as_i64() {
182                    if base_id == layer_id {
183                        children.push(PImgLayer::new(layer, layers)?);
184                    }
185                }
186            }
187        }
188        Ok(Self {
189            data,
190            layer_id,
191            layer_type,
192            name,
193            left,
194            top,
195            width,
196            height,
197            opacity,
198            visible,
199            type_,
200            children,
201        })
202    }
203
204    fn len(&self) -> usize {
205        1 + self.children.iter().map(|c| c.len()).sum::<usize>()
206    }
207
208    fn load_img(&self, img: &PImg) -> Result<ImageData> {
209        if self.layer_type == 2 || self.layer_type == 1 {
210            anyhow::bail!("Group layers do not have image data");
211        }
212        if self.layer_id == -1 {
213            // Generate a empty image
214            Ok(ImageData {
215                width: self.width,
216                height: self.height,
217                color_type: ImageColorType::Rgba,
218                depth: 8,
219                data: vec![0u8; (self.width * self.height * 4) as usize],
220            })
221        } else {
222            let tlg = img.load_img(self.layer_id).map_err(|e| {
223                anyhow::anyhow!("Failed to load image for layer_id {}: {}", self.layer_id, e)
224            })?;
225            let mut img = ImageData {
226                width: tlg.width,
227                height: tlg.height,
228                color_type: match tlg.color {
229                    TlgColorType::Bgr24 => ImageColorType::Bgr,
230                    TlgColorType::Bgra32 => ImageColorType::Bgra,
231                    TlgColorType::Grayscale8 => ImageColorType::Grayscale,
232                },
233                depth: 8,
234                data: tlg.data.clone(),
235            };
236            convert_to_rgba(&mut img)?;
237            Ok(img)
238        }
239    }
240
241    fn save_to_psd(&self, img: &PImg, psd: &mut PsdWriter, base: &mut ImageData) -> Result<()> {
242        if self.children.is_empty() {
243            let img = self.load_img(img)?;
244            let mut visible = self.visible;
245            if !self.data["diff_id"].is_none() {
246                visible = false; // Diff layers are always hide by default
247            }
248            if visible {
249                draw_on_img_with_opacity(base, &img, self.left, self.top, self.opacity)?;
250            }
251            let layer_name_source_setting = LayerNameSourceSetting {
252                id: self.layer_id as i32,
253            };
254            let mut packed = Vec::new();
255            layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
256            let additional_info = vec![AdditionalLayerInfo {
257                signature: *IMAGE_RESOURCE_SIGNATURE,
258                key: *LAYER_NAME_SOURCE_SETTING_KEY,
259                data: packed,
260            }];
261            let option = PsdLayerOption {
262                visible,
263                opacity: self.opacity,
264                additional_info,
265            };
266            psd.add_layer(self.name, self.left, self.top, img, Some(option))?;
267        } else {
268            psd.add_layer_group_end()?;
269            if self.layer_type == 0 {
270                let img = self.load_img(img)?;
271                let visible = self.visible;
272                if visible {
273                    draw_on_img_with_opacity(base, &img, self.left, self.top, self.opacity)?;
274                }
275                let layer_name_source_setting = LayerNameSourceSetting {
276                    id: self.layer_id as i32,
277                };
278                let mut packed = Vec::new();
279                layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
280                let additional_info = vec![AdditionalLayerInfo {
281                    signature: *IMAGE_RESOURCE_SIGNATURE,
282                    key: *LAYER_NAME_SOURCE_SETTING_KEY,
283                    data: packed,
284                }];
285                let option = PsdLayerOption {
286                    visible,
287                    opacity: self.opacity,
288                    additional_info,
289                };
290                psd.add_layer(self.name, self.left, self.top, img, Some(option))?;
291            }
292            for child in &self.children {
293                child.save_to_psd(img, psd, base)?;
294            }
295            let layer_name_source_setting = LayerNameSourceSetting {
296                id: self.layer_id as i32,
297            };
298            let mut packed = Vec::new();
299            layer_name_source_setting.pack(&mut packed, true, Encoding::Utf8, &None)?;
300            let additional_info = vec![AdditionalLayerInfo {
301                signature: *IMAGE_RESOURCE_SIGNATURE,
302                key: *LAYER_NAME_SOURCE_SETTING_KEY,
303                data: packed,
304            }];
305            let option = if self.layer_type == 0 {
306                Some(PsdLayerOption {
307                    additional_info,
308                    ..Default::default()
309                })
310            } else {
311                Some(PsdLayerOption {
312                    visible: self.visible,
313                    opacity: self.opacity,
314                    additional_info,
315                })
316            };
317            psd.add_layer_group(self.name, self.layer_type == 2, option)?;
318        }
319        Ok(())
320    }
321}
322
323#[derive(Debug)]
324struct PImgLayerRoot<'a> {
325    layers: Vec<PImgLayer<'a>>,
326}
327
328impl<'a> PImgLayerRoot<'a> {
329    pub fn new(layers: &'a PsbValueFixed) -> Result<Self> {
330        let mut root_layers = Vec::new();
331        for layer in layers.members() {
332            if layer["group_layer_id"].is_none() && layer["diff_id"].is_none() {
333                root_layers.push(PImgLayer::new(layer, layers)?);
334            }
335        }
336        Ok(Self {
337            layers: root_layers,
338        })
339    }
340
341    fn len(&self) -> usize {
342        self.layers.iter().map(|l| l.len()).sum()
343    }
344
345    fn save_to_psd(&self, img: &PImg, psd: &mut PsdWriter, base: &mut ImageData) -> Result<()> {
346        for layer in &self.layers {
347            layer.save_to_psd(img, psd, base)?;
348        }
349        Ok(())
350    }
351}
352
353#[derive(Debug)]
354/// Emote PImg Script
355pub struct PImg {
356    psb: VirtualPsbFixed,
357    overlay: Option<bool>,
358    psd: bool,
359    psd_compress: bool,
360    zlib_compression_level: u32,
361}
362
363impl PImg {
364    /// Create a new PImg script
365    ///
366    /// * `reader` - The reader containing the PImg script data
367    /// * `filename` - The name of the file
368    /// * `config` - Extra configuration options
369    pub fn new<R: Read + Seek>(reader: R, _filename: &str, config: &ExtraConfig) -> Result<Self> {
370        let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed();
371        Ok(Self {
372            psb,
373            overlay: config.emote_pimg_overlay,
374            psd: config.emote_pimg_psd,
375            psd_compress: config.psd_compress,
376            zlib_compression_level: config.zlib_compression_level,
377        })
378    }
379
380    fn load_img(&self, layer_id: i64) -> Result<Tlg> {
381        let layer_id = layer_id as usize;
382        let psb = self.psb.root();
383        let reference = &psb[format!("{layer_id}.tlg")];
384        let resource_id = reference
385            .resource_id()
386            .ok_or_else(|| anyhow::anyhow!("Layer {layer_id} does not have a resource ID"))?
387            as usize;
388        if resource_id >= self.psb.resources().len() {
389            return Err(anyhow::anyhow!(
390                "Resource ID {resource_id} for layer {layer_id} is out of bounds"
391            ));
392        }
393        let resource = &self.psb.resources()[resource_id];
394        Ok(load_tlg(MemReaderRef::new(&resource))?)
395    }
396}
397
398impl Script for PImg {
399    fn default_output_script_type(&self) -> OutputScriptType {
400        OutputScriptType::Custom
401    }
402
403    fn is_output_supported(&self, output: OutputScriptType) -> bool {
404        matches!(output, OutputScriptType::Custom)
405    }
406
407    fn custom_output_extension<'a>(&'a self) -> &'a str {
408        "psd"
409    }
410
411    fn default_format_type(&self) -> FormatOptions {
412        FormatOptions::None
413    }
414
415    fn is_image(&self) -> bool {
416        !self.psd
417    }
418
419    fn is_multi_image(&self) -> bool {
420        !self.psd
421    }
422
423    fn export_multi_image<'a>(
424        &'a self,
425    ) -> Result<Box<dyn Iterator<Item = Result<ImageDataWithName>> + 'a>> {
426        let psb = self.psb.root();
427        let overlay = self.overlay.unwrap_or_else(|| {
428            psb["layers"]
429                .members()
430                .all(|layer| layer["group_layer_id"].is_none())
431        });
432        if !overlay {
433            return Ok(Box::new(PImgIter2 {
434                pimg: self,
435                layers: psb.iter(),
436            }));
437        }
438        let width = psb["width"]
439            .as_u32()
440            .ok_or(anyhow::anyhow!("missing width"))?;
441        let height = psb["height"]
442            .as_u32()
443            .ok_or(anyhow::anyhow!("missing height"))?;
444        if !psb["layers"].is_list() {
445            return Err(anyhow::anyhow!("layers is not a list"));
446        }
447        if psb["layers"].len() == 0 {
448            return Ok(Box::new(std::iter::empty()));
449        }
450        let mut bases = HashMap::new();
451        for i in psb["layers"].members() {
452            if !i["diff_id"].is_none() {
453                continue; // Skip layers with diff_id
454            }
455            let layer_id = i["layer_id"]
456                .as_i64()
457                .ok_or(anyhow::anyhow!("missing layer_id"))?;
458            let top = i["top"].as_u32().ok_or(anyhow::anyhow!("missing top"))?;
459            let left = i["left"].as_u32().ok_or(anyhow::anyhow!("missing left"))?;
460            let opacity = i["opacity"]
461                .as_u8()
462                .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid opacity"))?;
463            bases.insert(layer_id, (self.load_img(layer_id)?, top, left, opacity));
464        }
465        Ok(Box::new(PImgIter {
466            pimg: self,
467            width,
468            height,
469            layers: psb["layers"].members(),
470            bases,
471        }))
472    }
473
474    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
475        let psb = self.psb.root();
476        let width = psb["width"]
477            .as_u32()
478            .ok_or(anyhow::anyhow!("missing width"))?;
479        let height = psb["height"]
480            .as_u32()
481            .ok_or(anyhow::anyhow!("missing height"))?;
482        let mut psd = PsdWriter::new(width, height, ImageColorType::Rgba, 8, encoding)?
483            .compress(self.psd_compress)
484            .zlib_compression_level(self.zlib_compression_level);
485        let mut base = ImageData {
486            width,
487            height,
488            color_type: ImageColorType::Rgba,
489            depth: 8,
490            data: vec![0u8; (width * height * 4) as usize],
491        };
492        let layers = PImgLayerRoot::new(&psb["layers"])?;
493        if layers.len() != psb["layers"].len() {
494            return Err(anyhow::anyhow!("Layer hierarchy is invalid"));
495        }
496        layers.save_to_psd(self, &mut psd, &mut base)?;
497        let file = std::fs::File::create(filename)?;
498        let mut writer = std::io::BufWriter::new(file);
499        psd.save(base, &mut writer)?;
500        Ok(())
501    }
502}
503
504struct PImgIter<'a> {
505    pimg: &'a PImg,
506    width: u32,
507    height: u32,
508    layers: ListIter<'a>,
509    bases: HashMap<i64, (Tlg, u32, u32, u8)>,
510}
511
512impl<'a> Iterator for PImgIter<'a> {
513    type Item = Result<ImageDataWithName>;
514
515    fn next(&mut self) -> Option<Self::Item> {
516        match self.layers.next() {
517            Some(layer) => {
518                let layer_id =
519                    try_option!(layer["layer_id"].as_i64().ok_or_else(|| {
520                        anyhow::anyhow!("Layer does not have a valid layer_id")
521                    }));
522                let layer_name = try_option!(
523                    layer["name"]
524                        .as_str()
525                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid name") })
526                );
527                let width = try_option!(
528                    layer["width"]
529                        .as_u32()
530                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid width") })
531                );
532                let height = try_option!(
533                    layer["height"]
534                        .as_u32()
535                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid height") })
536                );
537                let top = try_option!(
538                    layer["top"]
539                        .as_u32()
540                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid top") })
541                );
542                let left = try_option!(
543                    layer["left"]
544                        .as_u32()
545                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid left") })
546                );
547                let opacity = try_option!(
548                    layer["opacity"]
549                        .as_u8()
550                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid opacity") })
551                );
552                if layer["diff_id"].is_none() {
553                    let base = &try_option!(self.bases.get(&layer_id).ok_or(anyhow::anyhow!(
554                        "Base image for layer_id {} not found",
555                        layer_id
556                    )))
557                    .0;
558                    let mut data = ImageData {
559                        width: self.width,
560                        height: self.height,
561                        color_type: match base.color {
562                            TlgColorType::Bgr24 => ImageColorType::Bgr,
563                            TlgColorType::Bgra32 => ImageColorType::Bgra,
564                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
565                        },
566                        depth: 8,
567                        data: base.data.clone(),
568                    };
569                    if opacity != 255 {
570                        try_option!(apply_opacity(&mut data, opacity));
571                    }
572                    if self.width != width || self.height != height || top != 0 || left != 0 {
573                        data =
574                            try_option!(draw_on_canvas(data, self.width, self.height, left, top));
575                    }
576                    return Some(Ok(ImageDataWithName {
577                        name: layer_name.to_string(),
578                        data,
579                    }));
580                } else {
581                    let diff_id =
582                        try_option!(layer["diff_id"].as_i64().ok_or_else(|| {
583                            anyhow::anyhow!("Layer does not have a valid diff_id")
584                        }));
585                    let (base, base_top, base_left, base_opacity) = try_option!(
586                        self.bases
587                            .get(&diff_id)
588                            .ok_or(anyhow::anyhow!("Base image layer {} not found", diff_id))
589                    );
590                    let diff = try_option!(self.pimg.load_img(layer_id));
591                    if base.color != diff.color {
592                        return Some(Err(anyhow::anyhow!(
593                            "Color type mismatch for layer_id {}: base color {:?}, diff color {:?}",
594                            layer_id,
595                            base.color,
596                            diff.color
597                        )));
598                    }
599                    let mut base_img = ImageData {
600                        width: base.width,
601                        height: base.height,
602                        color_type: match base.color {
603                            TlgColorType::Bgr24 => ImageColorType::Bgr,
604                            TlgColorType::Bgra32 => ImageColorType::Bgra,
605                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
606                        },
607                        depth: 8,
608                        data: base.data.clone(),
609                    };
610                    if base.width != self.width
611                        || base.height != self.height
612                        || *base_top != 0
613                        || *base_left != 0
614                    {
615                        base_img = try_option!(draw_on_canvas(
616                            base_img,
617                            self.width,
618                            self.height,
619                            *base_left,
620                            *base_top
621                        ));
622                    }
623                    if *base_opacity != 255 {
624                        try_option!(apply_opacity(&mut base_img, *base_opacity));
625                    }
626                    let diff = ImageData {
627                        width: diff.width,
628                        height: diff.height,
629                        color_type: match diff.color {
630                            TlgColorType::Bgr24 => ImageColorType::Bgr,
631                            TlgColorType::Bgra32 => ImageColorType::Bgra,
632                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
633                        },
634                        depth: 8,
635                        data: diff.data.clone(),
636                    };
637                    try_option!(draw_on_img_with_opacity(
638                        &mut base_img,
639                        &diff,
640                        left,
641                        top,
642                        opacity
643                    ));
644                    Some(Ok(ImageDataWithName {
645                        name: layer_name.to_string(),
646                        data: base_img,
647                    }))
648                }
649            }
650            None => None,
651        }
652    }
653}
654
655struct PImgIter2<'a> {
656    pimg: &'a PImg,
657    layers: ObjectIter<'a>,
658}
659
660impl<'a> Iterator for PImgIter2<'a> {
661    type Item = Result<ImageDataWithName>;
662
663    fn next(&mut self) -> Option<Self::Item> {
664        match self.layers.next() {
665            Some((k, v)) => {
666                if !k.ends_with(".tlg") {
667                    return self.next();
668                }
669                let resource_id = try_option!(
670                    v.resource_id()
671                        .ok_or_else(|| anyhow::anyhow!("Layer {} does not have a resource ID", k))
672                ) as usize;
673                let name = k.trim_end_matches(".tlg").to_string();
674                if resource_id >= self.pimg.psb.resources().len() {
675                    return Some(Err(anyhow::anyhow!(
676                        "Resource ID {} for layer {} is out of bounds",
677                        resource_id,
678                        k
679                    )));
680                }
681                let resource = &self.pimg.psb.resources()[resource_id];
682                let tlg = try_option!(load_tlg(MemReaderRef::new(&resource)));
683                Some(Ok(ImageDataWithName {
684                    name,
685                    data: ImageData {
686                        width: tlg.width,
687                        height: tlg.height,
688                        color_type: match tlg.color {
689                            TlgColorType::Bgr24 => ImageColorType::Bgr,
690                            TlgColorType::Bgra32 => ImageColorType::Bgra,
691                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
692                        },
693                        depth: 8,
694                        data: tlg.data.clone(),
695                    },
696                }))
697            }
698            None => None,
699        }
700    }
701}